module psram #(
    parameter CLK_FREQ = 40000000
) (
    input clk,
    input resetn,

    output psram_clk,
    inout [3:0] psram_dio,
    output reg psram_cs,

    input [31:0] mem_addr,
    input [31:0] mem_wdata,
    input [3:0] mem_wstrb,
    output reg [31:0] mem_rdata,
    output reg mem_ready,
    input mem_valid

);
localparam qpi_idle  = 4'h0;
localparam qpi_cmd   = 4'h1;
localparam qpi_addr  = 4'h2;
localparam qpi_wait  = 4'h3;
localparam qpi_read  = 4'h4;
localparam qpi_write = 4'h5;
localparam qpi_finish = 4'h6;
wire [3:0] io_mode;//为0为输入
wire [3:0] psram_i;
wire [3:0] psram_o;

assign psram_i = psram_dio;

assign psram_dio[0] = io_mode[0] ? psram_o[0] : 1'bz;
assign psram_dio[1] = io_mode[1] ? psram_o[1] : 1'bz;
assign psram_dio[2] = io_mode[2] ? psram_o[2] : 1'bz;
assign psram_dio[3] = io_mode[3] ? psram_o[3] : 1'bz;


assign io_mode = (psram_state == qpi_cmd) ? 4'b0001 : 
                (psram_state == qpi_addr) ? 4'b1111 : 
                (psram_state == qpi_write) ? 4'b1111 : 4'b0000;
reg [7:0] shift_reg_o;
assign psram_o[0] = (psram_state == qpi_cmd) ? shift_reg_o[7] : shift_reg_o[4];
assign psram_o[1] = shift_reg_o[5];
assign psram_o[2] = shift_reg_o[6];
assign psram_o[3] = shift_reg_o[7];

reg start;
reg psram_rw;
reg psram_finish;
reg [3:0] psram_state;

reg clk_control;
assign psram_clk = clk_control&clk;
reg clk_control_pos;
reg [2:0] data_count;
reg [23:0] addr;
reg [23:0] addr_o;
reg [2:0] data_size;

reg [3:0] qpi_cnt;
reg [2:0] qpi_delay_cnt;
reg [1:0] addr_count;
reg [31:0] psram_wdata;
reg [31:0] psram_wdata_o;
reg [31:0] psram_rdata;
reg [7:0] shift_reg_i;
always @(negedge clk) begin
    if(!resetn) begin
        clk_control <= 1'b0;
    end begin
        clk_control <= clk_control_pos;
    end
end
always @(posedge clk or negedge resetn)begin
    if(!resetn)begin
        clk_control_pos <= 1'b0;
        psram_state <= qpi_idle;
        shift_reg_o <= 8'hff;
        addr_count <= 2'd0;
        qpi_delay_cnt <= 3'd0;
    end
    else begin
        case(psram_state)
            qpi_idle:begin
                clk_control_pos <= 1'b0;
                qpi_cnt <= 4'd0;
                qpi_delay_cnt <= 3'd0;
                data_count <= 3'd0;
                if(start == 1'b1)begin
                    psram_state <= qpi_cmd;
                    clk_control_pos <= 1'b1;
                    if(psram_rw == 1'b1)begin//写数据
                        shift_reg_o <= 8'h38;
                    end
                    else begin
                        shift_reg_o <= 8'hEB;
                    end
                end
                else begin
                    psram_state <= qpi_idle;
                end
            end
            qpi_cmd:begin
                qpi_cnt <= qpi_cnt + 1'b1;
                psram_state <= qpi_cmd;
                if(qpi_cnt != 4'h7)begin
                    shift_reg_o <= {shift_reg_o[6:0],1'b1};
                end
                else if(qpi_cnt == 4'd7)begin
                    qpi_cnt <= 4'd0;
                    addr_o <= {addr[15:0],8'hff};
                    shift_reg_o <= addr[23:16];
                    addr_count <= 2'd0;
                    psram_state <= qpi_addr;
                end
            end
            qpi_addr:begin
                psram_state <= qpi_addr;
                qpi_cnt <= qpi_cnt + 1'b1;
                if(qpi_cnt == 4'd1)begin
                    qpi_cnt <= 4'd0;
                    addr_o <= {addr_o[15:0],8'hff};
                    shift_reg_o <= addr_o[23:16];
                    if(addr_count != 2'd2)begin
                        addr_count <= addr_count + 1'b1;
                    end
                    else begin
                        data_count <= 3'd0;
                        if(psram_rw)begin
                            psram_state <= qpi_write;
                            qpi_cnt <= 4'd0;
                            shift_reg_o <= psram_wdata[7:0];
                            psram_wdata_o <= {8'hff,psram_wdata[31:8]};
                        end
                        else begin
                            qpi_delay_cnt <= 3'd0;
                            psram_state <= qpi_wait;
                            shift_reg_o <= 8'hff;
                        end
                    end
                end
                else begin
                    shift_reg_o <= {shift_reg_o[3:0],4'hf};
                end
            end
            qpi_wait:begin
                psram_state <= qpi_wait;
                qpi_delay_cnt <= qpi_delay_cnt + 1'b1;
                if(qpi_delay_cnt == 3'd5)begin
                    qpi_delay_cnt <= 3'd0;
                    qpi_cnt <= 4'd0;
                    data_count <= 3'd0;
                    psram_state <= qpi_read;
                end
            end
            qpi_read:begin
                psram_state <= qpi_read;
                shift_reg_i <= {shift_reg_i[3:0],psram_i[3:0]};                
                qpi_cnt <= qpi_cnt + 1'b1;
                if(qpi_cnt == 4'd1)begin
                    qpi_cnt <= 4'd0;
                    data_count <= data_count + 1'b1;
                    if(data_count == (data_size - 1'b1))begin
                        psram_state <= qpi_finish;                        
                        clk_control_pos <= 1'b0;
                    end
                end
                else begin
                    psram_rdata <= {shift_reg_i[7:0],psram_rdata[31:8]};
                end
            end
            qpi_write:begin
                psram_state <= qpi_write;
                qpi_cnt <= qpi_cnt + 1'b1;
                if(qpi_cnt == 4'd1)begin
                    qpi_cnt <= 4'd0;
                    shift_reg_o <= psram_wdata_o[7:0];
                    psram_wdata_o <= {8'hff,psram_wdata_o[31:8]};
                    data_count <= data_count + 1'b1;
                    if(data_count == (data_size - 1'b1))begin
                        psram_state <= qpi_finish;
                        clk_control_pos <= 1'b0;
                    end
                end
                else begin
                    shift_reg_o <= {shift_reg_o[3:0],4'hf};
                end
            end
            qpi_finish:begin
                mem_rdata <= {shift_reg_i[7:0],psram_rdata[31:8]};
                if(mem_valid == 1'b0)begin
                    psram_state <= qpi_idle;
                end
                else begin
                    psram_state <= qpi_finish;
                end
            end
        endcase
    end
end
reg mem_valid_r0;
always@(posedge clk) mem_valid_r0 <= mem_valid&resetn;
wire mem_valid_neg;
assign mem_valid_neg = (~mem_valid) & mem_valid_r0;
always @(posedge clk or negedge resetn) begin
    if(!resetn) begin
        mem_ready <= 1'b0;
        start <= 1'b0;
        psram_cs <= 1'b1;
        addr <= 24'd0;
        data_size <= 3'd0;
    end
    else if(mem_valid&&(mem_ready == 1'b0)) begin
        mem_ready <= 1'b0;
        psram_cs <= 1'b0;
        start <= 1'b1;
        
        if(psram_state == qpi_finish) begin
            mem_ready <= 1'b1;
            start <= 1'b0;
            psram_cs <= 1'b1;
        end

        if(mem_wstrb == 4'b0000)begin//读数据 spi:0xEB qpi:3字节地址 6个时钟周期 读取数据
            psram_rw <= 1'b0;
            data_size <= 3'd4;
            addr <= {mem_addr[23:2],2'b00};
        end
        else begin//写数据 spi:0x38 qpi:3字节地址 写入数据
            psram_rw <= 1'b1;
            if(mem_wstrb == 4'b0001)begin
                psram_wdata <= {4{mem_wdata[7:0]}};
                data_size <= 3'd1;
                addr <= {mem_addr[23:2],2'b00};
            end
            else if(mem_wstrb == 4'b0010)begin
                psram_wdata <= {4{mem_wdata[15:8]}};
                data_size <= 3'd1;
                addr <= {mem_addr[23:2],2'b01};
            end
            else if(mem_wstrb == 4'b0100)begin
                psram_wdata <= {4{mem_wdata[23:16]}};
                data_size <= 3'd1;
                addr <= {mem_addr[23:2],2'b10};
            end
            else if(mem_wstrb == 4'b1000)begin
                psram_wdata <= {4{mem_wdata[31:24]}};
                data_size <= 3'd1;
                addr <= {mem_addr[23:2],2'b11};
            end
            else if(mem_wstrb == 4'b0011)begin
                psram_wdata <= {2{mem_wdata[15:0]}};
                data_size <= 3'd2;
                addr <= {mem_addr[23:2],2'b00};
            end
            else if(mem_wstrb == 4'b1100)begin
                psram_wdata <= {2{mem_wdata[31:16]}};
                data_size <= 3'd2;
                addr <= {mem_addr[23:2],2'b10};
            end
            else begin
                psram_wdata <= mem_wdata;
                data_size <= 3'd4;
                addr <= {mem_addr[23:2],2'b00};
            end
        end
    end
    else if(mem_valid_neg) begin
        mem_ready <= 1'b0;
    end
    else begin
        start <= 1'b0;
        psram_cs <= 1'b1;
        mem_ready <= mem_ready;
    end
end
endmodule